from typing import TypedDict

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph,START,END
from  rich.console import Console
console = Console()

class State(TypedDict):
    input: str
    user_feedback: str

def step_1(state):
    print("-----step_1-----")
    pass
def human_feedback(state):
    print("-----human_feedback-----")
    pass
def step_3(state):
    print("-----step_3-----")
    pass

# 初始化图
builder = StateGraph(State)

# 创建节点
builder.add_node("step_1", step_1)
builder.add_node("human_feedback", human_feedback)
builder.add_node("step_3", step_3)

# 创建边
builder.add_edge(START, "step_1")
builder.add_edge("step_1", "human_feedback")
builder.add_edge("human_feedback", "step_3")
builder.add_edge("step_3", END)

# 初始化memory
memory = MemorySaver()

# compile 是传入checkpointer作为memory的缓存
# 然后指定断点，这样一来graph运行，就会停留在human_feedback之前来徐闻用户
graph = builder.compile(checkpointer=memory,interrupt_before=["human_feedback"])

initial_input = {"input": "你好"}
session = {"configurable": {"thread_id": "1"}}

for event in graph.stream(initial_input, session,stream_mode=["values"]):
    console.print(event)

user_input = input("Tell me you want to update the state:")
graph.update_state(session,{"user_feedback": user_input},as_node="human_feedback")

console.print("-----State after update-----",style="white on blue")

console.print(graph.get_state(session))

# graph.get_state(session).next
for event in graph.stream(None, session,stream_mode=["values"]):
    console.print(event)

